#ifndef _7DOF_MODLE_H_
#define _7DOF_MODLE_H_

#include <pinocchio/parsers/urdf.hpp> // 引入 URDF 解析器
#include <pinocchio/algorithm/kinematics.hpp>
#include <pinocchio/algorithm/jacobian.hpp>
#include <pinocchio/algorithm/frames.hpp>
#include <pinocchio/algorithm/joint-configuration.hpp>
#include <pinocchio/multibody/model.hpp>
#include <pinocchio/algorithm/rnea.hpp>

#include "ros/ros.h"

#include <Eigen/Dense>
#include <math.h>
#include <stdio.h>
using namespace std;
using namespace Eigen;

class UR7DOFARM
{

public:
    // 改进DH
    float a[7] = {0}, d[7] = {0}, alpha[7] = {0}, offset[7] = {0};
    // 标准DH
    float ja[7] = {0}, jd[7] = {0}, jalpha[7] = {0}, joffset[7] = {0};

    // 改进 DH
    void SetDH(float set_a[2], float set_d[2], float set_alpha[2])
    {
        for (size_t i = 0; i < 2; i++)
        {
            a[i] = set_a[i];
            d[i] = set_d[i];
            alpha[i] = set_alpha[i];
        }
    }

    // 改进 DH
    void SetDH_SD(float set_a[2], float set_d[2], float set_alpha[2])
    {
        for (size_t i = 0; i < 2; i++)
        {
            ja[i] = set_a[i];
            jd[i] = set_d[i];
            jalpha[i] = set_alpha[i];
        }
    }

    VectorXd Get_Tor(
        const VectorXd &q, const VectorXd &dq, const VectorXd &ddq,
        pinocchio::Model &model, pinocchio::Data &data);

    VectorXd CalculatePositionImpedanceTorque(
        const VectorXd &joint_pos,
        const VectorXd &joint_vel,
        const Eigen::Vector3d &desired_pos,
        const Eigen::Matrix3d &Kp,
        const Eigen::Matrix3d &Kd,
        pinocchio::Model &model,
        pinocchio::Data &data);
};

/**
 * @brief 计算机械臂的重力补偿力矩（逆动力学计算）
 *
 * @param[in] q     关节位置向量 (7x1, 单位: rad)
 * @param[in] dq    关节速度向量 (7x1, 单位: rad/s)
 * @param[in] ddq   关节加速度向量 (7x1, 单位: rad/s²)
 * @param[in] model Pinocchio机器人模型（必须已正确初始化）
 * @param[in] data  Pinocchio计算数据（需与model匹配）
 *
 * @return VectorXd 各关节重力补偿力矩 (7x1, 单位: N·m)
 */
VectorXd UR7DOFARM::Get_Tor(
    const VectorXd &q, const VectorXd &dq, const VectorXd &ddq,
    pinocchio::Model &model, pinocchio::Data &data)
{
    // 1. 设置重力方向（X轴方向）
    model.gravity.linear() << 9.81, 0, 0;

    // 2. 检查模型维度是否匹配
    assert(model.nq == 7 && "Model dimension doesn't match input size!");

    // 3. 直接计算逆动力学（无需类型转换）
    VectorXd tau = pinocchio::rnea(model, data, q, dq, ddq);

    // 4. 直接返回double类型
    return tau;
}

/**
 * @brief 计算七轴机械臂的纯位置阻抗控制力矩
 * @param joint_pos 当前关节位置 (7x1 VectorXd)
 * @param joint_vel 当前关节速度 (7x1 VectorXd)
 * @param desired_pos 期望末端位置 (3x1 Vector3d) [m]
 * @param Kp 位置刚度矩阵 (3x3 Matrix3d) [N/m]
 * @param Kd 阻尼矩阵 (3x3 Matrix3d) [N·s/m]
 * @param model Pinocchio模型
 * @param data Pinocchio数据
 * @return 阻抗控制力矩 (7x1 VectorXd)
 */
VectorXd UR7DOFARM::CalculatePositionImpedanceTorque(
    const VectorXd &joint_pos,
    const VectorXd &joint_vel,
    const Eigen::Vector3d &desired_pos,
    const Eigen::Matrix3d &Kp,
    const Eigen::Matrix3d &Kd,
    pinocchio::Model &model,
    pinocchio::Data &data)
{
    // 1. 计算当前末端位置
    pinocchio::forwardKinematics(model, data, joint_pos);
    int end_effector_id = model.getFrameId("link7");
    Eigen::Vector3d current_pos = data.oMf[end_effector_id].translation();

    ROS_INFO("pos: %.2f %.2f %.2f",
             current_pos.x(),current_pos.y(),current_pos.z());

    // 2. 计算位置误差
    Eigen::Vector3d position_error = desired_pos - current_pos;

    // 3. 计算雅可比矩阵（仅需线性部分）
    MatrixXd J = MatrixXd::Zero(6, model.nv);
    pinocchio::computeFrameJacobian(
        model, data, joint_pos, end_effector_id,
        pinocchio::LOCAL_WORLD_ALIGNED, J);
    MatrixXd J_linear = J.topRows<3>();

    // 4. 计算速度误差
    Eigen::Vector3d velocity_error = -J_linear * joint_vel;

    // 5. 计算笛卡尔空间阻抗力
    Eigen::Vector3d impedance_force = Kp * position_error + Kd * velocity_error;

    // 6. 转换为关节空间力矩
    return J_linear.transpose() * impedance_force;
}

#endif
